# 引入

import torch
import os
import glob
import random
import csv
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
from PIL import Image
from torchvision.transforms.transforms import CenterCrop, RandomRotation


class PokemonData(Dataset):
    # root 为数据根目录位置 ； resize ： 调整图片大小 ； mode ： 数据的类型（training，test)
    def __init__(self, root, resize, mode):
        super(PokemonData, self).__init__()

        self.root = root
        self.resize = resize

        self.name2label = {}
        for name in sorted(os.listdir(os.path.join(root))):
            if not os.path.isdir(os.path.join(root, name)):
                continue
            self.name2label[name] = len(self.name2label.keys())
        print(self.name2label)
        # image , label
        # 需要保存image_path而不不是image本身,如果把图片全部保存起来会炸
        self.images, self.labels = self.load_csv('images.csv')

        # if mode == 'train':  # 70%
        #     self.images = self.images[:int(0.7*len(self.images))]
        #     self.labels = self.labels[:int(0.7*len(self.labels))]
        # elif mode == 'val':  # 20%
        #     self.images = self.images[int(
        #         0.6*len(self.images)):int(0.8*len(self.images))]
        #     self.labels = self.labels[int(
        #         0.6*len(self.labels)):int(0.8*len(self.labels))]
        # else:  # 20%
        #     self.images = self.images[int(0.7*len(self.images)):]
        #     self.labels = self.labels[int(0.7*len(self.labels)):]

    # 返回labels和images两个列表

    def load_csv(self, filename):
        # 若csv文件不存在，则进入if进行创建
        if not os.path.exists(os.path.join(self.root, filename)):
            imageswrit = []
            for name in self.name2label.keys():
                # 保存到images这个列表中去
                # 保存内容为每张图片的地址（其中便包含label信息）
                imageswrit += glob.glob(os.path.join(self.root, name, '*.png'))

            # 共1165张图片如“pokeman/bulbasaur/00000047.png”
            print(len(imageswrit), imageswrit)
            # 打开文件，保存images
            # 写入文件
            with open(os.path.join(self.root, filename), mode='w', newline='')as f:
                writer = csv.writer(f)
                for img in imageswrit:  # pokeman/bulbasaur/00000047.png
                    name = img.split(os.sep)[-2]
                    label = self.name2label[name]
                    writer.writerow([img, label])
                # csv 文件可以得到【路径，label】对
        # read from csv
        # 打开刚才写入的文件
        images, labels = [], []
        with open(os.path.join(self.root, filename)) as f:
            reader = csv.reader(f)
            for row in reader:
                img, label = row
                label = int(label)

                images.append(img)
                labels.append(label)
        assert len(images) == len(labels)
        return images, labels

    def __len__(self):
        return len(self.images)

    def denormalize(self, x_hat):
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        # x_hat是如何得到的（归一化）
        # x_hat=(x-mean)/std
        # x=x_hat*std+mean
        # x_shape=[c,h,w]
        # mean_shape:[3]
        # 应当将mean转化成[3,1,1]
        mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
        std = torch.tensor(std).unsqueeze(1).unsqueeze(1)
        # 使用unsqueeze，给后面添加两个维度,1指的是位置
        x = x_hat*std+mean
        return x

    def __getitem__(self, index):
        img, label = self.images[index], self.labels[index]
        # lambda 可以自定义转化的方法，然后传入Compose作参数
        tf = transforms.Compose([

            # lambda x:Image.open(x).convert('RGB'),
            # transforms.Resize((self.resize, self.resize)),
            # transforms.ToTensor()

            lambda x:Image.open(x).convert('RGB'),
            # transforms.Resize((int(self.resize*1.25), int(self.resize*1.25))),
            # transforms.RandomRotation(15),  # 旋转会加大网络学习的难度,亦可以提升网络的准确性
            transforms.CenterCrop(self.resize),  # 裁减至符合大小（中心裁减）
            # 映射    #ToTensor是指把PIL.Image(RGB) 或者numpy.ndarray(H x W x C) 从0到255的值映射到0到1的范围内，并转化成Tensor格式。
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
                                 0.229, 0.224, 0.225])  # 归一化  # 统计imagenet上的所有图片所得到的参数，Normalize函数将图片的颜色值分布在0左右（-1～1）
        ])  # means为均值，std为方差
        img = tf(img)
        label = torch.tensor(label)
        return img, label


def showTensorTopic(a: int):  # 将tensor转化成pic的方法

    db = PokemonData('train', 200, 'train')
    BATCH_SIZE = 20
    loader = DataLoader(db, batch_size=BATCH_SIZE,
                        shuffle=True, num_workers=8)  # 加载一个batch的图片   # num_workers 为同时进行多线程提取图片

    fig = plt.figure()
    # 一页展示一个batch的图片
    for a, b in loader:
        for i in range(0, BATCH_SIZE-1):
            x, y = a[i], b[i]
            print('sample: ', x.shape, y.shape, y)
            # 取消归一化
            #x = db.denormalize(x)
            x = x.swapaxes(0, 1)
            x = x.swapaxes(1, 2)
            plt.subplot(3, 3, i+1)
            plt.tight_layout()
            plt.imshow(x)
            plt.title("Ground Truth: {}".format(y))
            plt.xticks([])
            plt.yticks([])
        plt.show()
    # 仅仅展示一张图片的方法
    #
    #x, y = db.__getitem__(a)
    # print('sample: ', x.shape, y.shape, y)
    # fig = plt.figure()
    # # x = db.denormalize(x)
    # x = x.swapaxes(0, 1)
    # x = x.swapaxes(1, 2)
    # plt.tight_layout()
    # plt.imshow(x)
    # plt.title("Ground Truth: {}".format(y))
    # plt.xticks([])
    # plt.yticks([])
    # plt.show()


def main():
    # db = PokemonData('pokeman', 224, 'train')
    # x, y = next(iter(db))  # iter()将db转化为迭代器，以可以使用next
    # print('sample: ', x.shape, y.shape, y)
    showTensorTopic(1)


if __name__ == '__main__':
    main()
